#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
检查天池竞赛数据格式

Author: BOSS (牛马)
Date: 2024-06-19
"""

import pandas as pd
import numpy as np
from pathlib import Path
import sys

def check_file(file_path):
    """检查单个数据文件"""
    print(f"\n{'='*60}")
    print(f"📊 检查文件: {file_path}")
    print(f"{'='*60}")
    
    try:
        # 读取数据
        df = pd.read_csv(file_path)
        
        # 基本信息
        print(f"✅ 文件读取成功")
        print(f"   数据形状: {df.shape}")
        print(f"   列名: {list(df.columns)}")
        
        # 显示前几行
        print(f"\n📋 前5行数据:")
        print(df.head())
        
        # 数据类型
        print(f"\n📊 数据类型:")
        for col in df.columns:
            print(f"   {col}: {df[col].dtype}")
        
        # 缺失值检查
        print(f"\n🔍 缺失值检查:")
        missing = df.isnull().sum()
        for col in df.columns:
            if missing[col] > 0:
                print(f"   {col}: {missing[col]} 个缺失值")
            else:
                print(f"   {col}: 无缺失值")
        
        # 如果有label列，显示标签分布
        if 'label' in df.columns:
            print(f"\n📈 标签分布:")
            label_counts = df['label'].value_counts().sort_index()
            for label, count in label_counts.items():
                print(f"   标签 {label}: {count} 个样本")
        
        # 文本长度统计
        text_columns = []
        for col in df.columns:
            if df[col].dtype == 'object' and col not in ['id', 'label']:
                text_columns.append(col)
        
        if text_columns:
            print(f"\n📝 文本长度统计:")
            for col in text_columns:
                lengths = df[col].astype(str).str.len()
                print(f"   {col}: 平均长度 {lengths.mean():.1f}, 最大长度 {lengths.max()}, 最小长度 {lengths.min()}")
        
        return df
        
    except Exception as e:
        print(f"❌ 读取文件失败: {e}")
        return None

def suggest_column_mapping(df):
    """建议列名映射"""
    print(f"\n💡 列名映射建议:")
    
    columns = df.columns.tolist()
    
    # 常见的列名模式
    mappings = {
        'id': ['id', 'ID', 'index', 'sample_id', 'idx'],
        'query1': ['query1', 'text1', 'sentence1', 'sent1', 'question1', 'q1', 'left', 'premise', 'text_a'],
        'query2': ['query2', 'text2', 'sentence2', 'sent2', 'question2', 'q2', 'right', 'hypothesis', 'text_b'],
        'label': ['label', 'target', 'class', 'similarity', 'match', 'is_duplicate']
    }
    
    suggested_mapping = {}
    
    for standard_name, possible_names in mappings.items():
        for col in columns:
            if col.lower() in [name.lower() for name in possible_names]:
                suggested_mapping[col] = standard_name
                print(f"   {col} -> {standard_name}")
                break
    
    # 检查未映射的列
    unmapped = [col for col in columns if col not in suggested_mapping]
    if unmapped:
        print(f"   未映射的列: {unmapped}")
    
    return suggested_mapping

def convert_to_standard_format(df, file_path):
    """转换为标准格式"""
    print(f"\n🔄 转换为标准格式...")
    
    # 建议映射
    mapping = suggest_column_mapping(df)
    
    # 创建标准格式的DataFrame
    standard_df = pd.DataFrame()
    
    # 映射已知列
    reverse_mapping = {v: k for k, v in mapping.items()}
    
    # ID列
    if 'id' in reverse_mapping:
        standard_df['id'] = df[reverse_mapping['id']]
    else:
        standard_df['id'] = range(len(df))
        print(f"   创建ID列")
    
    # query1列
    if 'query1' in reverse_mapping:
        standard_df['query1'] = df[reverse_mapping['query1']]
        print(f"   映射 {reverse_mapping['query1']} -> query1")
    else:
        # 尝试找到第一个文本列
        text_cols = [col for col in df.columns if df[col].dtype == 'object' and col not in mapping]
        if text_cols:
            standard_df['query1'] = df[text_cols[0]]
            print(f"   推测 {text_cols[0]} -> query1")
    
    # query2列
    if 'query2' in reverse_mapping:
        standard_df['query2'] = df[reverse_mapping['query2']]
        print(f"   映射 {reverse_mapping['query2']} -> query2")
    else:
        # 尝试找到第二个文本列
        text_cols = [col for col in df.columns if df[col].dtype == 'object' and col not in mapping]
        if len(text_cols) > 1:
            standard_df['query2'] = df[text_cols[1]]
            print(f"   推测 {text_cols[1]} -> query2")
    
    # label列（如果存在）
    if 'label' in reverse_mapping:
        standard_df['label'] = df[reverse_mapping['label']]
        print(f"   映射 {reverse_mapping['label']} -> label")
        
        # 检查标签格式
        unique_labels = standard_df['label'].unique()
        print(f"   标签值: {sorted(unique_labels)}")
        
        # 如果标签不是0/1格式，尝试转换
        if set(unique_labels) == {-1, 1}:
            standard_df['label'] = standard_df['label'].map({-1: 0, 1: 1})
            print(f"   标签已转换: -1/1 -> 0/1")
        elif not set(unique_labels).issubset({0, 1}):
            print(f"   ⚠️ 标签格式可能需要手动处理")
    
    # 保存转换后的文件
    output_dir = Path("data/raw")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    if 'train' in str(file_path).lower() or 'label' in standard_df.columns:
        output_path = output_dir / "train.csv"
    else:
        output_path = output_dir / "test.csv"
    
    standard_df.to_csv(output_path, index=False)
    print(f"✅ 标准格式文件已保存: {output_path}")
    
    return standard_df

def main():
    """主函数"""
    print("🔍 天池竞赛数据格式检查工具")
    print("Author: BOSS (牛马)")
    print("=" * 60)
    
    # 查找数据文件
    search_paths = [
        "data/raw/*.csv",
        "*.csv",
        "data/*.csv",
        "dataset/*.csv"
    ]
    
    found_files = []
    for pattern in search_paths:
        found_files.extend(Path().glob(pattern))
    
    if len(sys.argv) > 1:
        # 手动指定文件
        specified_files = [Path(arg) for arg in sys.argv[1:] if Path(arg).exists()]
        found_files.extend(specified_files)
    
    if not found_files:
        print("❌ 未找到CSV文件")
        print("\n📝 请确保数据文件在以下位置之一:")
        print("   - data/raw/train.csv")
        print("   - data/raw/test.csv")
        print("   - 当前目录下的CSV文件")
        print("\n💡 或者手动指定文件:")
        print("   python check_data.py your_file.csv")
        return
    
    print(f"✅ 找到 {len(found_files)} 个CSV文件")
    
    # 检查每个文件
    for file_path in found_files:
        df = check_file(file_path)
        if df is not None:
            # 询问是否转换格式
            print(f"\n❓ 是否将此文件转换为标准格式？")
            print(f"   文件: {file_path}")
            
            # 自动转换（在实际使用中可以改为交互式）
            convert_to_standard_format(df, file_path)
    
    print(f"\n🎉 数据检查完成！")
    print(f"📝 下一步:")
    print(f"   1. 检查 data/raw/ 目录下的标准格式文件")
    print(f"   2. 运行 python demo_standalone.py 开始训练")

if __name__ == "__main__":
    main()
